import argparse
import re
from string import strip

# macros
DATA_WIDTH = 32
MAX_WORDS = 2048
IMMEDIATE_LENGTH = 16
INST_2REG_IMMEDIATE = set(["addi", "andi", "ori", "hi", "beq", "bne"])
INST_2REG_IMMEDIATE_PAREN = set(["lw", "sw", "jal"])
INST_ALU = set(['and','or','xor','nand','nor','nxor','add','sub','eq','lt','le','ne','ge','gt'])
INST_PSEUDO = set(["subi", "not", "br", "blt", "ble", "bgt", "bge", "call", "jmp", "ret"])
NOOP = '0000000000000000'
PRIMARY_OPCODES = {
    'andi': '000001',
    'ori':  '000010',
    'hi':   '000011',
    'addi': '000100',
    'lw':   '010000',
    'sw':   '010001',
    'jal':  '100000',
    'beq':  '100001',
    'bne':  '100101'}
SECONDARY_OPCODES = {
    'and':  '001001',
    'or':   '001010',
    'xor':  '001011',
    'nand': '001101',
    'nor':  '001110',
    'nxor': '001111',
    'add':  '010000',
    'sub':  '010001',
    'eq':   '100001',
    'lt':   '100010',
    'le':   '100011',
    'ne':   '100101',
    'ge':   '100110',
    'gt':   '100111'}
REGISTERS = {
    'r0':   '00000',
    'r1':   '00001',
    'r2':   '00010',
    'r3':   '00011',
    'r4':   '00100',
    'r5':   '00101',
    'r6':   '00110',
    'r7':   '00111',
    'r8':   '01000',
    'r9':   '01001',
    'r10':   '01010',
    'r11':   '01011',
    'r12':   '01100',
    'r13':   '01101',
    'r14':   '01110',
    'r15':   '01111',
    'r16':   '10000',
    'r17':   '10001',
    'r18':   '10010',
    'r19':   '10011',
    'r20':   '10100',
    'r21':   '10101',
    'r22':   '10110',
    'r23':   '10111',
    'r24':   '11000',
    'r25':   '11001',
    'r26':   '11010',
    'r27':   '11011',
    'r28':   '11100',
    'r29':   '11101',
    'r30':   '11110',
    'r31':   '11111',
    'a0':    '00001',
    'a1':    '00010',
    'a2':    '00011',
    'a3':    '00100',
    't0':   '01000',
    't1':   '01001',
    't2':   '01010',
    't3':   '01011',
    't4':   '01100',
    't5':   '01101',
    't6':   '01110',
    't7':   '01111',
    's0':   '10000',
    's1':   '10001',
    's2':   '10010',
    's3':   '10011',
    's4':   '10100',
    's5':   '10101',
    's6':   '10110',
    's7':   '10111',
    'gp':   '11100',
    'fp':   '11101',
    'sp':   '11110',
    'ra':   '11111',
    'zero': '00000'}

# functions
def fatal(message):
    print message
    exit(1)

def strtoint(string):
    hex_pat = re.compile(r"0x([a-f0-9]+)", re.IGNORECASE)
    decimal_pat = re.compile(r"([-0-9]+)", re.IGNORECASE)
    hexval = hex_pat.search(string)
    decimalval = decimal_pat.search(string)
    if hexval:
        # hack to handle 32-bit values
        num_bits = 16
        if (len(hexval.group(1)) == 8): num_bits = 32
        num = int(hexval.group(1), 16)
        if num > 2**(num_bits-1)-1:
            num = num-2**(num_bits)
    elif decimalval:
        num = int(decimalval.group(1))
    return num

# this version of the function assumes you want an immediate
def sxtimm(num):
    return sxtimm_helper(num, IMMEDIATE_LENGTH)

def sxtimm_helper(num, length):
    if abs(num) > 2**(length)-1:
        fatal('immediate out of bounds')
    if num < 0:
        # find 2's comp, convert to binary, strip '0b'
        twoscomp = (2**length) + num
        binstr = bin(twoscomp)[2:]
        # sign extend
        for i in range(len(binstr), length):
            binstr = '1%s' % (binstr)
    else:
        binstr = bin(num)[2:]
        # sign extend
        for i in range(len(binstr), length):
            binstr = '0%s' % (binstr)
    return binstr

# handle arguments
parser = argparse.ArgumentParser(description="Converts assembly to machine code")
parser.add_argument("input_file", help="input assembly file")
parser.add_argument("-o", "--output-file", help="output machine code file")
parser.add_argument("-q", "--quiet", action="store_true")
parser.add_argument("-n", "--interleave-noops", type=int, default=0, help="number of noops to put between instructions")
options = parser.parse_args()

output = ''

# print header stuff
output += "WIDTH=%s;" % (DATA_WIDTH) + '\n'
output += "DEPTH=2048;" + '\n'
output += "ADDRESS_RADIX=HEX;" + '\n'
output += "DATA_RADIX=HEX;" + '\n'
output += "CONTENT BEGIN" + '\n'

# read in the a32 file
f = open(options.input_file, 'r')
source = [line.strip().lower() for line in f.readlines()]
# remove useless lines (blank or comment)
useless_pat = re.compile(r'^\s*(?:;.*)?$')
source = filter(lambda str:not useless_pat.match(str), source)
# close file
f.close()

# define regex for recognizing syntax
pseudoop_pat = re.compile(r"^\.(NAME|ORIG|WORD)\s+(\w+)\s*(?:=\s*(\w+))?", re.IGNORECASE)
#instruction_pat = re.compile("^(\w+)\s+", re.IGNORECASE)
instruction_pat = re.compile("^(\w+)", re.IGNORECASE)
label_pat = re.compile("^(\w+):", re.IGNORECASE)

three_op_reg_pat = re.compile(r'\w+\s+(\w+)\s*,\s*(\w+)\s*,\s*(\w+)\s*(?:;.*)?$')
three_op_pat = re.compile(r"^\w+\s+(\w+)\s*,\s*(\w+)\s*,\s*((?:-?[0-9]+)|(?:0x[0-9a-f]+))\s*(?:;.*)?$", re.IGNORECASE)
three_op_label_pat = re.compile(r"^\w+\s+(\w+)\s*,\s*(\w+)\s*,\s*([a-z]\w*)", re.IGNORECASE)
three_op_paren_pat = re.compile(r"^\w+\s+(?P<reg1>\w+)\s*,\s*(?:(?P<labelimm>[a-z]\w*)|(?P<litimm>(?:0x[0-9a-f]+)|(?:-?[0-9]+)))\((?P<reg2>\w+)\)", re.IGNORECASE)
two_op_pat = re.compile(r"^\w+\s+(\w+)\s*,\s*(\w+)\s*(?:;.*)?$")
two_op_paren_pat = re.compile(r"^\w+\s+(\w+\(\w+\))\s*(?:;.*)?$", re.IGNORECASE)
one_label_pat = re.compile(r"\w+\s+([a-z]\w*)\s*(?:;.*)?$")

# prep data structure for machine code
machine_code=[]
for i in range(MAX_WORDS):
    machine_code.append(('',''))

### PASS ONE
## build symbol table
lc=0
symbol_table=dict()
# config table holds tuples (location, label)
config_table={'references': [], 'offsets':[], 'references32':[]}
#for line in source:
i = 0
while i < len(source):
    line = source[i]
    increment_i = True
    increment_lc = True
    pseudoop = pseudoop_pat.match(line)
    label = label_pat.match(line)
    instruction = instruction_pat.match(line)
    if pseudoop:
        if pseudoop.group(1) == 'name':
            symbol_table[pseudoop.group(2)] = strtoint(pseudoop.group(3))
            increment_lc = False
        elif pseudoop.group(1) == 'orig':
            target = strtoint(pseudoop.group(2))/4
            lc = target
            increment_lc = False
        elif pseudoop.group(1) == 'word':
            # cases for hex, dec, and label
            if re.match(r'0x[0-9a-fA-F]+', pseudoop.group(2), re.IGNORECASE):
                val = int(pseudoop.group(2), 16)
                machine_code[lc] = (sxtimm_helper(val, 32), line)
            elif re.match(r'-?[0-9]+', pseudoop.group(2), re.IGNORECASE):
                val = int(pseudoop.group(2))
                machine_code[lc] = (sxtimm_helper(val, 32), line)
            else:
                config_table['references32'].append((lc, pseudoop.group(2)))
                machine_code[lc] = ('', line)
        else:
            fatal('invalid pseudoop')
    elif label:
        symbol_table[label.group(1)] = lc
        increment_lc = False
    elif instruction:
        machine_line = ''
        inst = instruction.group(1)
        if inst in INST_2REG_IMMEDIATE:
            # instructions of the format:
            # <inst> <reg>,<reg>,<immediate>
            ops = three_op_pat.match(line)
            ops_label = three_op_label_pat.match(line)
            # construct machine code up to immediate
            # figure out how to handle immediate
            if ops:
                # in this case, immediate is literal
                machine_line = '%s%s%s' % (PRIMARY_OPCODES[inst], REGISTERS[ops.group(2)], REGISTERS[ops.group(1)])
                immediate = strtoint(ops.group(3))
                machine_line += sxtimm(immediate)
            elif ops_label:
                # in this case, immediate is a label, so dereference it
                # if label not in symbol table, make a note config table
                machine_line = '%s%s%s' % (PRIMARY_OPCODES[inst], REGISTERS[ops_label.group(2)], REGISTERS[ops_label.group(1)])
                immediate = ops_label.group(3)
		# hack to fix test2: label should be byte address, not word
                if immediate in symbol_table:
                    immediate = symbol_table[immediate]
		    if inst == "addi":
		        immediate = immediate * 4
                    machine_line += sxtimm(immediate)
                else:
                    config_table['references'].append((lc, immediate))
            else:
                fatal('malformed instruction: %s' % (line))
        elif inst in INST_2REG_IMMEDIATE_PAREN:
            # instructions of the format:
            # <inst> <reg>,<immediate>(<reg>)
            ops = three_op_paren_pat.match(line)
            if ops:
                machine_line = '%s%s%s' % (PRIMARY_OPCODES[inst], REGISTERS[ops.group('reg2')], REGISTERS[ops.group('reg1')])
                # figure out what type of immediate we're given
                if ops.group('labelimm'):
                    # immediate is a label
                    immediate = ops.group('labelimm')
                    if immediate in symbol_table:
                        immediate = symbol_table[immediate]
                        machine_line += sxtimm(immediate)
                    else:
                        config_table['references'].append((lc, immediate))
                elif ops.group('litimm'):
                    immediate = ops.group('litimm')
                    immediate = strtoint(immediate)
                    machine_line += sxtimm(immediate)
                else:
                    fatal('malformed immediate operand: %s' % (line))
            else:
                fatal('malformed instruction: %s' % (line))
        elif inst in INST_ALU:
            # instructions of the format:
            # <inst> <dest reg>,<reg>,<reg>
            ops = three_op_reg_pat.match(line)
            if ops:
                machine_line = '000000%s%s%s00000%s' % (REGISTERS[ops.group(2)], REGISTERS[ops.group(3)], REGISTERS[ops.group(1)], SECONDARY_OPCODES[inst])
            else:
                fatal('malformed instruction: %s' % (line))
        elif inst in INST_PSEUDO:
            # instructions that must be translated to another instruction
            if inst == "subi":
                ops = three_op_pat.match(line)
                immediate = strtoint(ops.group(3))
                immediate *= -1
                translation = "addi %s,%s,%s" % (ops.group(1), ops.group(2), immediate)
            elif inst == 'not':
                ops = two_op_pat.match(line)
                translation = "nor %s,%s,%s" % (ops.group(1), ops.group(2), ops.group(2))
            elif inst == 'br':
                ops = one_label_pat.match(line)
                translation = "beq zero,zero,%s" % (ops.group(1))
            elif inst == 'blt':
                # this instruction and the next few require two lines
                ops = three_op_label_pat.match(line)
                translation = 'lt r6,%s,%s' % (ops.group(1), ops.group(2))
                source.insert(i,translation)
                i += 1
                translation = 'bne r6,zero,%s' % (ops.group(3))
            elif inst == 'ble':
                ops = three_op_label_pat.match(line)
                translation = 'le r6,%s,%s' % (ops.group(1), ops.group(2))
                source.insert(i,translation)
                i += 1
                translation = 'bne r6,zero,%s' % (ops.group(3))
            elif inst == 'bge':
                ops = three_op_label_pat.match(line)
                translation = 'ge r6,%s,%s' % (ops.group(1), ops.group(2))
                source.insert(i,translation)
                i += 1
                translation = 'bne r6,zero,%s' % (ops.group(3))
            elif inst == 'bgt':
                ops = three_op_label_pat.match(line)
                translation = 'gt r6,%s,%s' % (ops.group(1), ops.group(2))
                source.insert(i,translation)
                i += 1
                translation = 'bne r6,zero,%s' % (ops.group(3))
            elif inst == 'call':
                ops = two_op_paren_pat.match(line)
                translation = "jal ra,%s" % (ops.group(1))
            elif inst == 'jmp':
                ops = two_op_paren_pat.match(line)
                translation = "jal r6,%s" % (ops.group(1))
            elif inst == 'ret':
                translation = "jal r6,0(ra)"
            source[i] = translation
            # this is a hack to make the loop iterate over the second added inst
            if inst in set(['blt','ble','bge','bgt']):
                i -= 1
            increment_lc = False
            increment_i = False
        machine_code[lc] = (machine_line, line)
        if increment_lc:
            for j in range(1,options.interleave_noops+1):
                machine_code[lc+j] = (NOOP, 'noop')
    if (increment_i): i += 1
    if (increment_lc): lc += 1+options.interleave_noops

### PASS TWO
## handle config table entries
for (marker, label) in config_table['references']:
    (machine_line, line) = machine_code[marker]
    value = symbol_table[label]
    if machine_line[0:6] == PRIMARY_OPCODES['addi']:
        print "addi!"
	value = value*4
	print value
    machine_line += sxtimm(value)
    machine_code[marker] = (machine_line, line)
for (marker, label) in config_table['references32']:
    (machine_line, line) = machine_code[marker]
    machine_line = sxtimm_helper(symbol_table[label], 32)
    machine_code[marker] = (machine_line, line)

### PASS THREE
## convert absolute addresses to relative/offset
# possible issue here if data starts with beq/bne opcode
for lc in range(len(machine_code)):
    (machine_line, line) = machine_code[lc]
    if machine_line[:6] == PRIMARY_OPCODES['beq'] or machine_line[:6] == PRIMARY_OPCODES['bne']:
        absolute = machine_line[16:]
        absolute = int(absolute, 2)
        offset = absolute - (lc + 1)
        offset = sxtimm(offset)
        machine_code[lc] = ("%s%s" % (machine_line[:16], offset), line)

### PASS FOUR
## output
in_empty_region = False
for lc in range(len(machine_code)):
    (machine_line, line) = machine_code[lc]
    if machine_line == '':
        if in_empty_region:
            empty_region_end = lc
        else:
            in_empty_region = True
            empty_region_start = lc
    else:
        if in_empty_region:
            in_empty_region = False
            output += "[%08x..%08x] : DEAD;" % (empty_region_start, empty_region_end) + '\n'
        output += "-- @ 0x%08x :    %s" % (4*lc, line) + '\n'
        output += "%08x : %08x;" % (lc, int(machine_line, 2)) + '\n'
if in_empty_region:
    output += "[%08x..%08x] : DEAD;" % (empty_region_start, lc) + '\n'

output += 'END;'
if options.output_file:
    f = open(options.output_file, 'w')
    f.write(output)
    f.close
elif not options.quiet:
    print output
